"""
And a python3 version:
def unpickle(file):
    import pickle
    with open(file, 'rb') as fo:
        dict = pickle.load(fo, encoding='bytes')
    return dict
Loaded in this way, each of the batch files contains a dictionary with the following elements:
data -- a 10000x3072 numpy array of uint8s. Each row of the array stores a 32x32 colour image. The first 1024 entries
contain the red channel values, the next 1024 the green, and the final 1024 the blue. The image is stored in row-major
order, so that the first 32 entries of the array are the red channel values of the first row of the image.
labels -- a list of 10000 numbers in the range 0-9. The number at index i indicates the label of the ith image in the
array data.

The dataset contains another file, called batches.meta. It too contains a Python dictionary object. It has the following
entries:
label_names -- a 10-element list which gives meaningful names to the numeric labels in the labels array described above.
For example, label_names[0] == "airplane", label_names[1] == "automobile", etc.
"""
import pickle
from python_ai.common.xcommon import *
import matplotlib.pyplot as plt

path = r'../../../../large_data/DL1/cifar-10-batches-py/'
with open(path + 'data_batch_1', 'br') as f:
    dict = pickle.load(f, encoding='bytes')
sep('dict keys')
print(dict.keys())
data = dict[b'data']
labels = dict[b'labels']
check_shape(data, 'data')
check_shape(labels, 'labels')

# label texts
with open(path + 'batches.meta', 'br') as f:
    dict = pickle.load(f, encoding='bytes')
sep('dict keys for meta')
print(dict.keys())
labels_txt = dict[b'label_names']

n_cls = len(np.unique(labels))
sep(f'n_cls: {n_cls}')

data = data.reshape([-1, 3, 32, 32])
data = data.transpose([0, 2, 3, 1])

spr = 4
spc = 6
spn = 0
plt.figure(figsize=[12, 8])

for i in range(spr * spc):
    img = data[i]
    label = labels[i]
    label = labels_txt[label].decode()  # ATTENTION bytes.decode()
    spn += 1
    plt.subplot(spr, spc, spn)
    plt.axis('off')
    plt.imshow(img)
    plt.title(label)

plt.show()

